{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spherical Fourier Neural Operators\n",
    "\n",
    "A simple notebook to showcase spherical Fourier Neural Operators\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import paddle\n",
    "import paddle.nn as nn\n",
    "from paddle.io import DataLoader\n",
    "from paddle import amp\n",
    "from paddle.optimizer.lr import OneCycleLR\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "from math import ceil, sqrt\n",
    "\n",
    "import time\n",
    "\n",
    "cmap='twilight_shifted'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enable_amp = False\n",
    "\n",
    "# set device\n",
    "device = paddle.device.set_device('gpu' if paddle.device.cuda.device_count() > 0 else 'cpu')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training data\n",
    "to train our geometric FNOs, we require training data. To this end let us prepare a Dataloader which computes results on the fly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset\n",
    "from paddle_harmonics.examples.sfno import PdeDataset\n",
    "\n",
    "# 1 hour prediction steps\n",
    "dt = 1*3600\n",
    "dt_solver = 150\n",
    "nsteps = dt//dt_solver\n",
    "dataset = PdeDataset(dt=dt, nsteps=nsteps, dims=(256, 512), device=device, normalize=True)\n",
    "# There is still an issue with parallel dataloading. Do NOT use it at the moment\n",
    "dataloader = DataLoader(dataset, places=paddle.CUDAPinnedPlace(), batch_size=4, shuffle=True, num_workers=0, persistent_workers=False)\n",
    "solver = dataset.solver.to(device)\n",
    "\n",
    "nlat = dataset.nlat\n",
    "nlon = dataset.nlon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paddle.seed(0)\n",
    "inp, tar = dataset[0]\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(inp[2], fig, vmax=3, vmin=-3)\n",
    "plt.title(\"input\")\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(tar[2], fig, vmax=3, vmin=-3)\n",
    "plt.title(\"target\")\n",
    "plt.colorbar(im)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the geometric Fourier Neural Operator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from paddle_harmonics.examples.sfno import SphericalFourierNeuralOperatorNet as SFNO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SFNO(spectral_transform='sht', operator_type='driscoll-healy', img_size=(nlat, nlon), grid=\"equiangular\",\n",
    "                 num_layers=4, scale_factor=3, embed_dim=16, big_skip=True, pos_embed=\"lat\", use_mlp=False, normalization_layer=\"none\").to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def l2loss_sphere(solver, prd, tar, relative=False, squared=True):\n",
    "    loss = solver.integrate_grid((prd - tar)**2, dimensionless=True).sum(axis=-1)\n",
    "    if relative:\n",
    "        loss = loss / solver.integrate_grid(tar**2, dimensionless=True).sum(axis=-1)\n",
    "    \n",
    "    if not squared:\n",
    "        loss = paddle.sqrt(loss)\n",
    "    loss = loss.mean()\n",
    "\n",
    "    return loss\n",
    "\n",
    "def spectral_l2loss_sphere(solver, prd, tar, relative=False, squared=True):\n",
    "    # compute coefficients\n",
    "    coeffs = paddle.as_real(solver.sht(prd - tar))\n",
    "    coeffs = coeffs[..., 0]**2 + coeffs[..., 1]**2\n",
    "    norm2 = coeffs[..., :, 0] + 2 * paddle.sum(coeffs[..., :, 1:], axis=-1)\n",
    "    loss = paddle.sum(norm2, axis=(-1,-2))\n",
    "\n",
    "    if relative:\n",
    "        tar_coeffs = paddle.as_real(solver.sht(tar))\n",
    "        tar_coeffs = tar_coeffs[..., 0]**2 + tar_coeffs[..., 1]**2\n",
    "        tar_norm2 = tar_coeffs[..., :, 0] + 2 * paddle.sum(tar_coeffs[..., :, 1:], axis=-1)\n",
    "        tar_norm2 = paddle.sum(tar_norm2, axis=(-1,-2))\n",
    "        loss = loss / tar_norm2\n",
    "\n",
    "    if not squared:\n",
    "        loss = paddle.sqrt(loss)\n",
    "    loss = loss.mean()\n",
    "\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training function\n",
    "def train_model(model, dataloader, optimizer, scheduler=None, nepochs=20, nfuture=0, num_examples=256, num_valid=8, loss_fn='l2'):\n",
    "\n",
    "    train_start = time.time()\n",
    "\n",
    "    for epoch in range(nepochs):\n",
    "\n",
    "        # time each epoch\n",
    "        epoch_start = time.time()\n",
    "\n",
    "        dataloader.dataset.set_initial_condition('random')\n",
    "        dataloader.dataset.set_num_examples(num_examples)\n",
    "\n",
    "        optimizer.clear_grad(set_to_zero=True)\n",
    "\n",
    "        # do the training\n",
    "        acc_loss = 0\n",
    "        model.train()\n",
    "\n",
    "        for inp, tar in dataloader:\n",
    "            with amp.auto_cast(enable=enable_amp):\n",
    "                prd = model(inp)\n",
    "                for _ in range(nfuture):\n",
    "                    prd = model(prd)\n",
    "                if loss_fn == 'l2':\n",
    "                    loss = l2loss_sphere(solver, prd, tar)\n",
    "                elif loss_fn == \"spectral l2\":\n",
    "                    loss = spectral_l2loss_sphere(solver, prd, tar)\n",
    "\n",
    "            acc_loss += loss.item() * inp.shape[0]\n",
    "\n",
    "            optimizer.clear_grad(set_to_zero=True)\n",
    "            # gscaler.scale(loss).backward()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            # gscaler.update()\n",
    "\n",
    "        if scheduler is not None:\n",
    "            scheduler.step()\n",
    "\n",
    "        acc_loss = acc_loss / len(dataloader.dataset)\n",
    "\n",
    "        dataloader.dataset.set_initial_condition('random')\n",
    "        dataloader.dataset.set_num_examples(num_valid)\n",
    "\n",
    "        # perform validation\n",
    "        valid_loss = 0\n",
    "        model.eval()\n",
    "        with paddle.no_grad():\n",
    "            for inp, tar in dataloader:\n",
    "                prd = model(inp)\n",
    "                for _ in range(nfuture):\n",
    "                    prd = model(prd)\n",
    "                loss = l2loss_sphere(solver, prd, tar, relative=True)\n",
    "\n",
    "                valid_loss += loss.item() * inp.shape[0]\n",
    "\n",
    "        valid_loss = valid_loss / len(dataloader.dataset)\n",
    "\n",
    "        epoch_time = time.time() - epoch_start\n",
    "\n",
    "        print(f'--------------------------------------------------------------------------------')\n",
    "        print(f'Epoch {epoch} summary:')\n",
    "        print(f'time taken: {epoch_time}')\n",
    "        print(f'accumulated training loss: {acc_loss}')\n",
    "        print(f'relative validation loss: {valid_loss}')\n",
    "\n",
    "    train_time = time.time() - train_start\n",
    "\n",
    "    print(f'--------------------------------------------------------------------------------')\n",
    "    print(f'done. Training took {train_time}.')\n",
    "    return valid_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set seed\n",
    "paddle.seed(333)\n",
    "\n",
    "optimizer = paddle.optimizer.Adam(parameters=model.parameters(), learning_rate=3E-3, weight_decay=0.0)\n",
    "gscaler = amp.GradScaler(enable=enable_amp)\n",
    "train_model(model, dataloader, optimizer, nepochs=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader.dataset.set_initial_condition('random')\n",
    "\n",
    "paddle.seed(0)\n",
    "\n",
    "with paddle.no_grad():\n",
    "    inp, tar = next(iter(dataloader))\n",
    "    out = model(inp).detach()\n",
    "\n",
    "s = 0; ch = 2\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(inp[s, ch], fig, projection='3d', title='input')\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(out[s, ch], fig, projection='3d', title='prediction')\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(tar[s, ch], fig, projection='3d', title='target')\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata((tar-out)[s, ch], fig, projection='3d', title='error')\n",
    "plt.colorbar(im)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
